#include "stft.h"

namespace STFT_SPACE{
    struct Complex {
        double R, I;
        Complex(const double& R = 0.0, const double& I = 0.0) : R(R), I(I) {}
        double abs(){ return sqrt(R*R + I*I); }
    };

    inline Complex operator+(const Complex &C1, const Complex &C2) { return Complex{C1.R + C2.R, C1.I + C2.I}; }
    inline Complex operator-(const Complex &C1, const Complex &C2) { return Complex{C1.R - C2.R, C1.I - C2.I}; }
    inline Complex operator*(const Complex &C1, const Complex &C2) { return Complex{C1.R * C2.R - C1.I * C2.I, C1.R * C2.I + C1.I * C2.R}; }
    inline Complex operator~(const Complex &C) { return Complex{C.R, -C.I}; }
    inline Complex &operator*=(Complex &C1, const Complex &C2) { return C1 = C1 * C2; }

    constexpr double Pi = 3.1415926535897932384626433832795;
    constexpr int Max_size = 1 << 18 | 5;
    Complex complex_buffer[Max_size];
    int size;
    int bitrev[Max_size];
    Complex w[Max_size];

    inline void init(int n)
    {
        for (size = 1; size < n; size <<= 1);
        for (int i = 0; i != size; ++i)
            bitrev[i] = bitrev[i >> 1] >> 1 | ((i & 1) * (size >> 1));
        size >>= 1;
        for (int i = 0; i != size; ++i)
            w[size + i] = {cos(Pi * i / size), sin(Pi * i / size)};
        for (int i = size - 1; i; --i)
            w[i] = w[i << 1];
        size <<= 1;
    }

    inline void DFT(Complex A[], int L)
    {
        int shift = __builtin_ctz(size) - __builtin_ctz(L);
        for (int i = 0; i != L; ++i)
            if (i < (bitrev[i] >> shift))
                std::swap(A[i], A[bitrev[i] >> shift]);
        for (int d = 1; d != L; d <<= 1)
            for (int i = 0; i != L; i += d << 1)
                for (int j = 0; j != d; ++j)
                {
                    Complex tmp = A[i + d + j] * w[d + j];
                    A[i + d + j] = A[i + j] - tmp, A[i + j] = A[i + j] + tmp;
                }
    }

    inline void Real_DFT(Complex A[], int L)
    {
        if (L == 1)
            return;
        static Complex tmp[Max_size];
        L >>= 1;
        for (int i = 0; i != L; ++i)
            tmp[i] = {A[i << 1].R, A[i << 1 | 1].R};
        DFT(tmp, L);
        tmp[L] = tmp[0];
        for (int i = 0; i != L; ++i)
        {
            Complex t1 = (tmp[i] + ~tmp[L - i]) * Complex{0.5, 0}, t2 = (tmp[i] - ~tmp[L - i]) * Complex{0, -0.5} * w[L + i];
            A[i] = t1 + t2, A[L + i] = t1 - t2;
        }
    }

    inline void Real_IDFT(Complex A[], int L)
    {
        if (L == 1)
            return;
        static Complex tmp[Max_size];
        A[L] = A[0];
        std::reverse(A + 1, A + L);
        L >>= 1;
        for (int i = 0; i != L; ++i)
        {
            Complex t1 = (A[i] + A[L + i]) * Complex{0.5, 0}, t2 = (A[i] - A[L + i]) * Complex{0, 0.5} * w[L + i];
            tmp[i] = t1 + t2;
        }
        DFT(tmp, L);
        for (int i = 0; i != L; ++i)
            A[i << 1] = {tmp[i].R / L, 0}, A[i << 1 | 1] = {tmp[i].I / L, 0};
    }
}
void init_rfft(int L) {
    STFT_SPACE::init(L);
}
void rfft(std::vector<double> & A, int L) {
    STFT_SPACE::Complex * data = STFT_SPACE::complex_buffer;
    for(int i = 0; i < L; ++ i) data[i] = {A[i],0.0};
    Real_DFT(data,L);
    for(int i = 0; i < L; ++ i) A[i] = data[i].abs();
}
void irfft(std::vector<double> & A, int L) {
    STFT_SPACE::Complex * data = STFT_SPACE::complex_buffer;
    for(int i = 0; i < L; ++ i) data[i] = {A[i],0.0};
    Real_IDFT(data,L);
    for(int i = 0; i < L; ++ i) A[i] = data[i].abs();
}
void STFT(sound & me, double timeStep, double windowLength, int windowType, std::vector<std::vector<double>> & SF) {
    int wlen = floor(windowLength * me.getfrequency());
    int inc = floor(timeStep * me.getfrequency());
    SF = me.enframe(wlen,inc,windowType);
    init_rfft(wlen);
    int nFFT = STFT_SPACE::size, halfnFFT = nFFT/2;

    for(auto & i : SF) {
        i.resize(nFFT,0);
        rfft(i, nFFT);
        for(int j = 0; j < i.size(); ++ j) {
            i[j] = abs(i[j] / nFFT);
        }
        i.resize(nFFT/2+1);
        for(int j = 0; j < i.size(); ++ j) {
            if(j != 0 && j != halfnFFT) i[j] *= 2.0;
            i[j] = 20*log10(i[j]);
        }
    }
}
